# -*- coding:utf-8 -*-

# @Time    : 2023/4/2 02:37
# @Author  : zengwenjia
# @Email   : zengwenjia@lingxi.ai
# @File    : handle_dialogue.py
# @Software: LLM_internal

# # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # #

import re
import copy
from util_tool import utils
import os


CONTEXT_HEINIU = """
你接下来要作为一名电话销售员，与用户进行电话沟通，用户之前领取过泰康的一个意外险，你需要根据用户的需求，向用户推荐升级成泰康的一款百万医疗险，最终达到销售目标。
"""

CONTEXT_SHUIDI = """
你接下来要作为一名电话销售员，与用户进行电话沟通，用户之前领取过水滴的重疾险的一个赠险,你需要根据用户的需求，向用户推荐水滴保的长险产品，最终达到销售目标。
"""

SYSTEM = """
你首先需要根据用户性别和用户姓氏，和用户核实身份
"""
# 定义对话的记录单元类，包括角色，内容
class DialogRecord:
    def __init__(self, dialogue_id, role, content):
        self.dialogue_id = dialogue_id
        self.role = role
        self.content = content
        self.info = ""



class DialogAnalysis:
    def __init__(self):
        pass

    def read_data(self):
        # 读取txt文件
        file = open("dialogue_data/heiniu_dialogue_0330.txt", "r", encoding="utf-8")
        lines = file.readlines()
        file.close()

        dialogue_list = []
        index = 0
        record_list = []
        for line in lines:
            # 通过：分割，第一个是角色，第二个是内容
            line = line.replace("：", ":")
            line = line.replace("\n", "").split(":")
            # 去除空格
            line = [l.strip() for l in line]
            # 去除空行
            if len(line) == 2:

                if line[0] == "sessionId":
                    dialogue_index = line[1]
                    continue
                if (len(record_list) == 0):
                    dialog_record = DialogRecord(dialogue_index, "系统", SYSTEM)
                    record_list.append(dialog_record)
                    if (line[0] == "用户"):
                        continue

                # 去掉line[1]中的方括[]及里面的内容
                pattern = r'\[[^\]]*\]'
                line[1] = re.sub(pattern, '', line[1])
                line[1] = line[1].replace(" ", "")
                dialog_record = DialogRecord(dialogue_index, line[0], line[1])
                record_list.append(dialog_record)
            elif len(line) < 2:
                if len(record_list) != 0:
                    dialogue_list.append(record_list)
                index += 1
                dialogue_index = index
                record_list = []
        return dialogue_list

    def _format_dialogue(self, instruction: str, response: str = ''):
        if not response:
            return f'\n\n### Input:\n{instruction}'
        if not instruction:
            return f'\n\n### Response:\n{response}'
        return f'\n\n### Input:\n{instruction}\n\n### Response:\n{response}'

    #将chat_dialogues转化成prompt格式
    def _convert_to_prompt(self, chat_dialogues):
        prompt = ''
        for chat_dialogue in chat_dialogues:
            prompt += self._format_dialogue(chat_dialogue.instruction, chat_dialogue.response)
        return prompt

    # 根据对话的list，转化成promt格式
    def convert_to_prompt(self, dialogue_list):
        instruction_data = []
        for dialogue_text in dialogue_list:
            response = ""
            instruction = ""
            chat_dialogues = []
            for index, record in enumerate(dialogue_text):
                if (record.role == "机席") or (record.role == "人席")  or (record.role == "坐席"):
                    if not instruction:
                        continue
                    response = record.content
                    dialogue = Dialogue(instruction=instruction, response="")
                    chat_dialogues_task = copy.deepcopy(chat_dialogues)
                    chat_dialogues_task.append(dialogue)
                    result_dict = {}
                    result_dict["instruction"] = CONTEXT_HEINIU
                    result_dict["input"] = self._convert_to_prompt(chat_dialogues_task)
                    result_dict["output"] = response
                    print(result_dict)
                    instruction_data.append(result_dict)


                    dialogue = Dialogue(instruction=instruction, response=response)
                    chat_dialogues.append(dialogue)
                    response = ""
                    instruction = ""
                elif (record.role == "用户") or (record.role == "系统"):
                    instruction = record.content

        utils.jdump(instruction_data, os.path.join("../data_set/dialogue_data", "heiniu_instruction_0330.json"))

    # 根据用户核身话术，生成用户信息
    def generate_user_info(self, path):
        dialogue_list = utils.jload(path)
        new_dialogue_list = []
        for dialogue in dialogue_list:
            instruction = dialogue["instruction"]
            input = dialogue["input"]
            output = dialogue["output"]
            identification_sentence = ""
            if re.search(r'您好~请问您是.*吗', input):
                identification_sentence = re.search(r'您好~请问您是.*吗', input).group()
            elif re.search(r'您好~请问您是.*吗', output):
                identification_sentence = re.search(r'您好~请问您是.*吗', output).group()
            if identification_sentence:
                family_name = identification_sentence[7:8]
                gender_name = identification_sentence[8:10]
                if gender_name == "先生":
                    gender = "男"
                elif gender_name == "女士":
                    gender = "女"
            instruction = instruction + ",你首先需要根据用户性别和用户姓氏，和用户核实身份"
            if identification_sentence:
                instruction = instruction + "，用户姓" + family_name + ",性别是" + gender
            instruction = instruction + "\n"

            input = input.replace("\n\n### Input:\n\n你首先需要根据用户性别和用户姓氏，和用户核实身份\n", "")
            new_dialogue = {}
            new_dialogue["instruction"] = instruction
            new_dialogue["input"] = input.replace("哦您不要挂电话马上为您服务。", "").replace(
                "哦您不要挂电话马上为您服务啊。", "").replace("嗯您不要挂电话马上为您服务。", "").replace(
                "哦您不要挂电话马上为您服务，", "").replace("哦您不要挂电话马上为您服。", "").replace(
                "哦您不要挂电话，马上为您服务。", "").replace("哦您不要挂电话马上为您。", "").replace(
                "哦您不要挂电话马上为您服务", "").replace("您先不要挂电话好马上为您服务。", "").replace(
                "嗯您不要挂电话，马上为您服务。", "").replace("哦您不要挂电话马上为。", "").replace("电话马上为您服务。","")
            new_dialogue["output"] = output.replace("哦您不要挂电话马上为您服务。", "").replace(
                "哦您不要挂电话马上为您服务啊。", "").replace("嗯您不要挂电话马上为您服务。", "").replace(
                "哦您不要挂电话马上为您服务，", "").replace("哦您不要挂电话马上为您服。", "").replace(
                "哦您不要挂电话，马上为您服务。", "").replace("哦您不要挂电话马上为您。", "").replace(
                "哦您不要挂电话马上为您服务", "").replace("您先不要挂电话好马上为您服务。", "").replace(
                "嗯您不要挂电话，马上为您服务。", "").replace("哦您不要挂电话马上为。", "").replace("电话马上为您服务。","")
            if new_dialogue["output"]:
                new_dialogue_list.append(new_dialogue)
        utils.jdump(new_dialogue_list, os.path.join("../data_set/dialogue_data", "heiniu_instruction_0330_modify.json"))





if __name__ == "__main__":
    dialog_analysis = DialogAnalysis()
    # dialogue_list = dialog_analysis.read_data()
    # dialog_analysis.convert_to_prompt(dialogue_list)
    dialog_analysis.generate_user_info("dialogue_data/heiniu_instruction_0330.json")

